import os
import cv2
import time
import copy
import numpy as np

def computeSIFT1V0(img1, img2):
    if len(img1.shape)==3 and len(img2.shape)==3:
        img1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
        img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)

    sift = cv2.SIFT_create()
    keypoints_roi, descriptor_roi = sift.detectAndCompute(img1, None)
    keypoints_img, descriptor_img = sift.detectAndCompute(img2, None)

    return keypoints_roi, descriptor_roi, keypoints_img, descriptor_img

def computeSIFT(img):
    if len(img.shape)==3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    image = copy.deepcopy(img)
    h, w = image.shape
    image = cv2.resize(image, (h//2, w//2))
    sift = cv2.SIFT_create()
    keypoints_roi, descriptor_roi = sift.detectAndCompute(image, None)
    return keypoints_roi, descriptor_roi

def matchKeypoint(img1, img2, keypoints_roi, descriptor_roi, keypoints_img, descriptor_img, factor=3, vis=False):
    max_dist, min_dist = 0, 1000
    matcher = cv2.FlannBasedMatcher()
    matches = matcher.match(descriptor_roi, descriptor_img)
    good_match = []

    for i, match in enumerate(matches):
        if match.distance < min_dist:
            min_dist = match.distance
        if match.distance > max_dist:
            max_dist = match.distance

    for i, match in enumerate(matches):
        if match.distance < factor*min_dist:
            good_match.append(match)
    good_match = tuple(good_match)
    if vis:
        image = cv2.drawMatches(img1, keypoints_roi, img2, keypoints_img, good_match, None)
        return good_match, image
    else:
        return good_match, None

def matchKeypointKNN(img1, img2, keypoints_roi, descriptor_roi, keypoints_img, descriptor_img, factor=0.65, vis=False):
    max_dist, min_dist = 0, 1000
    matcher = cv2.BFMatcher()
    matches = matcher.knnMatch(descriptor_roi, descriptor_img, k=2)
    matches = sorted(matches, key=lambda x: x[0].distance/x[1].distance)
    good_match = []

    for m, n in matches:
        if m.distance < factor*n.distance:
            good_match.append(m)

    good_match = tuple(good_match)

    if vis:
        image = cv2.drawMatches(img1, keypoints_roi, img2, keypoints_img, good_match, None)
        return good_match, image
    else:
        return good_match, None

def findH(keypoints_roi, keypoints_img, good_match):
    roi = np.float32([keypoints_roi[m.queryIdx].pt for m in good_match]).reshape(-1, 1, 2)
    img = np.float32([keypoints_img[m.trainIdx].pt for m in good_match]).reshape(-1, 1, 2)
    H, status = cv2.findHomography(img, roi, cv2.RANSAC, ransacReprojThreshold=10)
    return H

def match(ref, tar, keypoints_roi, descriptor_roi, points, vis=False):
    img_save = tar
    shape = tar.shape
    # start = time.time()
    keypoints_img, descriptor_img = computeSIFT(tar)
    # finish1 = time.time()
    good_match, image = matchKeypointKNN(ref, tar, keypoints_roi, descriptor_roi, keypoints_img, descriptor_img, factor=0.65)
    H = findH(keypoints_roi, keypoints_img, good_match)
    # finish2 = time.time()

    colour = []
    for i, point in enumerate(points):
        img = img_save
        if i+1 > 5:   # 判断矩形个数，大于5个跳出
            break
        xmin, ymin, xmax, ymax = point
        #print(xmin, ymin, xmax, ymax)
        if vis:
            cv2.rectangle(ref, (xmin, ymax), (xmax, ymin), (0, 0, 255), thickness=3, lineType=4)
            cv2.imwrite('ref.jpg', ref)
        coordinate = [[xmin, ymin, 1], [xmax, ymax, 1]]

        # 矩阵相乘
        match_point = [np.dot(H, coordinate[0]), np.dot(H, coordinate[1])]
        xmin = (int(match_point[0][0]) if int(match_point[0][0]) > 0 else 0)
        ymin = (int(match_point[0][1]) if int(match_point[0][1]) > 0 else 0)
        xmax = (int(match_point[1][0]) if int(match_point[1][0]) < shape[1] else shape[1])
        ymax = (int(match_point[1][1]) if int(match_point[1][1]) < shape[0] else shape[0])
        # print(xmin, ymin, xmax, ymax)

        if vis:
            cv2.rectangle(tar, (xmin, ymax), (xmax, ymin), (0, 0, 255), thickness=2, lineType=4)  # 画矩形
            cv2.imwrite('tar.jpg', tar)

        img_ = img[ymin:ymax, xmin:xmax, :]   # 抠图
        #cv2.imwrite('c{}.jpg'.format(i), img_)

        colour.append(img_)

    # finish3 = time.time()
    # print('Match file: find keypoints: ', finish1-start)
    # print('Match file: find H matrix: ', finish2-finish1)
    # print('Match file: projection: ', finish3-finish2)
    #imgs = np.hstack([ref, tar])  # 将ref和tar显示在一张图上
    #cv2.imwrite('result.jpg', imgs)

    return colour



# if __name__ == '__main__':
#     img1 = cv2.imread('./171035.10.jpg')
#     img2 = cv2.imread('./171044.10.jpg')
#     y, x = img2.shape[0:2]
#     img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
#     img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
#     txt_path = r'F:\Postgraduate documents\project\door\chengdu\sample\sample1\rzj\2\black_2\C50.FR-L.10.txt'
#     points = [[]]
#     number = 0
#     # with open(txt_path, 'r') as f:
#     #     for line in f.readlines():
#     #         line = line.strip('\n')
#     #         list = line.split(' ')
#     #         xmin = (float(list[1])-float(list[3])) * x
#     #         xmax = (float(list[1])+float(list[3])) * x
#     #         ymin = (float(list[2])-float(list[4])) * y
#     #         ymax = (float(list[2])+float(list[4])) * y
#     #         print(xmin, ymin, xmax, ymax)
#     #         points[number] = [xmin, ymin, xmax, ymax]
#     #         number += 1
#
#     points = [[4888, 2199, 5088, 2392], [5186, 2361, 5386, 2554]]
#     colour = match_divide(img1, img2, points)
#     print(len(colour))
#
#     # keypoints_roi, descriptor_roi, keypoints_img, descriptor_img = computeSIFT(img1, img2)
#     # good_match, image = matchKeypointKNN(img1, img2, keypoints_roi, descriptor_roi, keypoints_img, descriptor_img, factor=0.65)
#     # cv2.imwrite('/media/Harddisk/Users/fan/PROJECT/CDPROJECT/TensorRT/Tools/169.254.7.13.jpg', image)
#     # print(good_match)
#     # H = findH(keypoints_roi, keypoints_img, good_match)
#     # # leftTop, leftBottom, rightTop, rightBottom = calcFourCorner(img2, H)
#     # image = showSimpleImg(img1, img2, H)
#     # cv2.imwrite('parliament_stich.jpg', image)
#     print('over')